# -*- coding: utf-8 -*-   
# @File   : evaluate.py         
# @Author : wjf
# @Time   : 2023/12/29 10:38
import pandas as pd


class EvaluateWarnResult(object):
    def __init__(self, predict_xlsx):
        self.predict_xlsx = predict_xlsx

    def get_v3_result(self):
        df = pd.read_excel(self.predict_xlsx)
        df['v3_result'] = [""] * df.shape[0]
        # predict
        condition_predict_truepredict_true1 = (df['common_black_flag'] == 1) & (df['model_flag'] == 1) & (df['super_black_flag'] == 1)
        condition_predict_truepredict_true2 = (df['common_black_flag'] == 1) & (df['model_flag'] == 1) & (df['super_black_flag'] == 0) & (df['super_white_flag'] == 0)
        condition_predict_truepredict_true3 = (df['common_black_flag'] == 0) & (df['score_model'] >= 0.5) & (df['start_white_flag'] == 0) & (df['end_black_flag'] == 1)
        condition_predict_truepredict_true4 = (df['common_black_flag'] == 0) & (df['score_model'] >= 0.5) & (df['start_white_flag'] == 0) & (df['end_black_flag'] == 0) & (
                df['end_white_flag'] == 0)

        condition_predict_true = (
                condition_predict_truepredict_true1 | condition_predict_truepredict_true2 | condition_predict_truepredict_true3 | condition_predict_truepredict_true4)
        condition_predict_false = ~condition_predict_true

        # label
        condition_label_true = (df['label'] == '1-红色') | (df['label'] == '2-橙色') | (df['label'] == '3-黄色') | (df['label'] == '4-蓝色')
        condition_label_false = ~condition_label_true

        tp = condition_label_true & condition_predict_true
        fp = condition_label_false & condition_predict_true
        tn = condition_label_false & condition_predict_false
        fn = condition_label_true & condition_predict_false

        df.loc[tp, 'v3_result'] = "tp"
        df.loc[fp, 'v3_result'] = "fp"
        df.loc[tn, 'v3_result'] = "tn"
        df.loc[fn, 'v3_result'] = "fn"

        df.to_excel(self.predict_xlsx.replace(".xlsx", "_eval_v3.xlsx"), index=False)

    def get_model_result(self):
        df = pd.read_excel(self.predict_xlsx)
        df['model_result'] = [""] * df.shape[0]

        # predict
        condition_predict_true = df['model_flag'] == 1
        condition_predict_false = df['model_flag'] == 0

        # label
        # label
        condition_label_true = (df['label'] == '1-红色') | (df['label'] == '2-橙色') | (df['label'] == '3-黄色') | (df['label'] == '4-蓝色')
        condition_label_false = ~condition_label_true

        tp = condition_label_true & condition_predict_true
        fp = condition_label_false & condition_predict_true
        tn = condition_label_false & condition_predict_false
        fn = condition_label_true & condition_predict_false

        df.loc[tp, 'model_result'] = "tp"
        df.loc[fp, 'model_result'] = "fp"
        df.loc[tn, 'model_result'] = "tn"
        df.loc[fn, 'model_result'] = "fn"

        df.to_excel(self.predict_xlsx.replace(".xlsx", "_eval_model.xlsx"), index=False)


if __name__ == "__main__":
    ewr = EvaluateWarnResult(predict_xlsx=r'result/test5000_pred_result_merge.xlsx')
    ewr.get_v3_result()
    ewr.get_model_result()
    pass
